package org.ztr.aiollama.controller;

import lombok.RequiredArgsConstructor;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.ztr.aiollama.repository.ChatHistoryRepository;
import reactor.core.publisher.Flux;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @Author: ZhaoTR
 * @Date: Created in 2025/9/20 16:52
 * @Description: AI-ollama 聊天控制器
 * @Version: 1.0
 */
//@RequiredArgsConstructor
@RestController
@RequestMapping("/ai")
public class ChatController {

    private static final Logger logger = LogManager.getLogger(ChatController.class);

    private final ChatClient chatClient;

    private final ChatHistoryRepository chatHistoryRepository;

    public ChatController(
            ChatClient chatClient,
            ChatHistoryRepository chatHistoryRepository
    ) {
        this.chatClient = chatClient;
        this.chatHistoryRepository = chatHistoryRepository;
    }

    @RequestMapping(value = "/chat", produces = "text/html;charset=utf-8")
    public Flux<String> chat(String prompt, String chatId) {

        chatHistoryRepository.save("chat", chatId);

        return chatClient.prompt()
                .user(prompt)
//                .call()
                .advisors(a -> a
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY,
                        10))
                .stream()
                .content();
    }

}
